import os
import csv
import datetime as dt
import tushare as ts#调取数据的包
import pandas as pd

from config import ENVIRONMENT_CONFIG

def csv_write(file, row, new=False):
    '''
    往csv文件写入一行数据
    :param file: 文件名
    :param row: 行数
    :param new: 是否是新建文件
    :return:
    '''
    if new:
        with open(file, 'w', newline='') as csvFile:
            writer = csv.writer(csvFile)
            writer.writerow(row)
    else:
        with open(file, 'a', newline='') as csvFile:
            writer = csv.writer(csvFile)
            writer.writerow(row)

def date_type_transfer1(date2):
    '''
    将%Y%m%d转化为%Y-%m-%d
    :param date2:
    :return:
    '''
    date1 = dt.datetime.strftime(dt.datetime.strptime(date2, '%Y%m%d'), '%Y-%m-%d')
    return date1

def date_type_transfer2(date1):
    '''
    将%Y-%m-%d转化为%Y%m%d
    :param date1:
    :return:
    '''
    date2 = dt.datetime.strftime(dt.datetime.strptime(date1, '%Y-%m-%d'), '%Y%m%d')
    return date2

def download_data(date):
    '''
    下载某一日的股票数据，写入文件
    :param date:
    :return:
    '''
    pro = ts.pro_api()
    df = pro.daily(trade_date=date_type_transfer2(date))
    df.to_csv(os.path.join(ENVIRONMENT_CONFIG['data_path'], '%s.csv' % date))

def download_trade_cal():#获取交易日历
    '''
    获取交易日历
    :return:
    '''
    pro = ts.pro_api()
    df = pro.query('trade_cal')#拿下日历
    df['cal_date'] = df['cal_date'].apply(date_type_transfer1)
    df[['cal_date', 'is_open']].to_csv(os.path.join(ENVIRONMENT_CONFIG['data_path'], 'trade_cal.csv'))

def get_trade_date(start_date, end_date):#选出需要交易日期的数据
    '''
    获取从start_date到end_date的数据
    :param start_date:
    :param end_date:
    :return:
    '''
    if not os.path.exists(os.path.join(ENVIRONMENT_CONFIG['data_path'], 'trade_cal.csv')):
        download_trade_cal()

    df = pd.read_csv(os.path.join(ENVIRONMENT_CONFIG['data_path'], 'trade_cal.csv'), index_col=0)
    df = df.query('cal_date >= "%s" and cal_date <= "%s" and is_open == 1' % (start_date, end_date))

    return df['cal_date'].tolist()

def get_data(date):
    '''
    只获取某一天的数据，并写入文件
    :param date:
    :return:
    '''
    if not os.path.exists(os.path.join(ENVIRONMENT_CONFIG['data_path'], '%s.csv' % date)):
        download_data(date)
    df = pd.read_csv(os.path.join(ENVIRONMENT_CONFIG['data_path'], '%s.csv' % date), index_col=0)
    return df[['ts_code', 'open', 'high', 'low', 'close', 'pre_close', 'vol', 'amount']].set_index('ts_code')

def get_stocklist():#筛选股票
    '''
    按照一定的条件选择股票信息，日期再20200101之前的股票
    :return: 股票代码列表
    '''
    pro = ts.pro_api()
    data = pro.stock_basic(exchange='', list_status='L', fields='ts_code,list_date')#l：正在上市
    return data.query('list_date < "20200101"')['ts_code'].tolist()